package sink

import (
	"context"
	"errors"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
	"go.uber.org/zap"
)

const SCHEMA_WS = "ws"

type wsSink struct {
	ctx             context.Context
	cancelWithCause context.CancelCauseFunc

	conn *websocket.Conn
}

func NewWsSink(url *url.URL) (zap.Sink, error) {
	wss := new(wsSink)
	wss.ctx, wss.cancelWithCause = context.WithCancelCause(context.Background())
	wss.attach(url)

	return wss, nil
}

func (w *wsSink) Close() error {
	w.detach(nil)
	return nil
}

func (w *wsSink) Sync() error {
	return nil
}

func (w *wsSink) Write(p []byte) (n int, err error) {
	select {
	case <-w.ctx.Done():
		err = errors.New("sink closed")
		w.detach(err)
	default:
		err = w.conn.WriteMessage(websocket.TextMessage, p)
		if err != nil {
			w.detach(err)
		}
	}
	return len(p), err
}

func (w *wsSink) attach(url *url.URL) error {
	conn, _, err := websocket.DefaultDialer.Dial(url.String(), nil)
	if err != nil {
		return err
	}
	w.conn = conn
	return nil
}

func (w *wsSink) detach(resultErr error) {
	w.cancelWithCause(resultErr)
	w.conn.WriteControl(
		websocket.CloseMessage,
		websocket.FormatCloseMessage(websocket.CloseNormalClosure, "goodbye"),
		time.Now().Add(5*time.Second),
	)
	w.conn.Close()
}
